import argparse
import psycopg2
import datetime
import pandas as pd
import datasketches

def get_conn():
    gs_conn = psycopg2.connect(
        host='yourHost',
        port='yourport',
        database='gauss_database_name',
        user='xxx',
        password='xxx'
    )
    print("连接数据库成功")
    return gs_conn

def getdf():
    print("开始读取高斯数据库实现datasketches功能")
    conn_py=get_conn()
    # 创建一个数据库游标
    cur = conn_py.cursor()
    # 执行查询语句
    cur.execute('select * from hjm.datasketches_t')
    #获取查询
    rows = cur.fetchall()
    # 转换为 DataFrame 对象
    pd_df = pd.DataFrame(rows, columns=[desc[0] for desc in cur.description])
    print("原始数据:",pd_df)

    # 创建Theta Sketch
    theta1 = datasketches.update_theta_sketch()
    # 查询第一组产品类别（假设类别ID为1）的产品ID并更新Theta Sketch
    cur.execute("SELECT id FROM hjm.datasketches_t WHERE product_id = 1")
    rows = cur.fetchall()
    for row in rows:
        product_id = row[0]
        theta1.update(product_id)

    # 将ThetaSketch转换为紧凑形式
    compact_theta_sketch = theta1.compact()
    # 将紧凑草图转换为字节串
    sketch_bytes = compact_theta_sketch.serialize()
    print("序列字符串:",sketch_bytes)
    #反序列字符串，转换为草图
    compact_theta_sketch_read = datasketches.compact_theta_sketch.deserialize(sketch_bytes)
    print("反序列信息：")
    print(compact_theta_sketch_read)
    # 获取基数估计值等操作
    estimate = compact_theta_sketch_read.get_estimate()
    print("ThetaSketch草图的基数估计值:", estimate)
    print("读取成功，且草图计算已生成")
    cur.close()
    conn_py.close()
    return sketch_bytes

def writedf(sketch_bytes):
    print("开始将草图信息写回高斯数据库")
    conn = get_conn()
    cur = conn.cursor()
    #将草图信息写回gaussdb
    sketch_info = {
        'sketch_type': 'ThetaSketch',
        'creation_time': datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
        'data_bytes': sketch_bytes
    }

    cur.execute("INSERT INTO hjm.datasketches_t_tmp (sketch_type, creation_time, data_bytes) VALUES (%s, %s, %s)",
                   (sketch_info['sketch_type'], sketch_info['creation_time'], sketch_info['data_bytes']))
    conn.commit()
    print("写入成功")
    conn.close()
    cur.close()
    print("数据库连接已关闭")

sketch_bytes = getdf()
writedf(sketch_bytes)